package org.ruoyi.chat.controller.tripartite;

import com.volcengine.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.volcengine.ark.runtime.model.images.generation.ImagesResponse;
import com.volcengine.ark.runtime.model.images.generation.ResponseFormat;
import com.volcengine.ark.runtime.service.ArkService;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.ruoyi.chat.domain.SeedObject;
import org.ruoyi.chat.domain.TaskResult;
import org.ruoyi.chat.util.RedisTaskUtil;
import org.ruoyi.common.oss.core.OssClient;
import org.springframework.web.bind.annotation.*;

import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@Api(tags = "火山接口")
@RestController
@RequestMapping("/seed")
@RequiredArgsConstructor
@Slf4j
public class SeedController {
    private final RedisTaskUtil redisTaskUtil;
    private final ExecutorService exector = Executors.newFixedThreadPool(5);

    @ApiOperation(value = "查询任务状态")
    @GetMapping("/status/{taskId}")
    public String getTaskStatus(@PathVariable String taskId) {
        TaskResult result = redisTaskUtil.getTaskResult(taskId);
        if (result == null) {
            return "任务ID不存在或已过期";
        }
        return result.getStatus();
    }

    @ApiOperation(value = "查询任务状态")
    @GetMapping("/result")
    public String getTaskResult(@RequestParam("taskId") String taskId) {
        TaskResult result = redisTaskUtil.getTaskResult(taskId);
        if (result == null) {
            return "任务ID不存在或已过期";
        }
        return result.getResult();
    }

    @ApiOperation(value = "火山seed融合图片")
    @PostMapping("/blend")
    public String blend(@RequestBody SeedObject seedObject) {
        String taskId = UUID.randomUUID().toString();
        redisTaskUtil.saveTaskResult(taskId, new TaskResult("PROCESSING", null, 0), 86400);

        exector.submit(()->{
            try{
                String resultUrl = processImageBlend(seedObject);
                redisTaskUtil.saveTaskResult(taskId, new TaskResult("COMPLETED", resultUrl, 100), 86400);
            } catch (Exception e) {
                // 异常时更新状态
                redisTaskUtil.saveTaskResult(taskId, new TaskResult("FAILED", e.getMessage(), 0), 86400);
            }
        });

        return taskId;
    }

    private String processImageBlend(SeedObject seedObject) throws Exception {
        String apiKey = "4686154b-9e1a-4759-92bd-98133ecf05c0";
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).apiKey(apiKey).build();

        // "保持图片中的位置关系，将图片中的家具自然地融合到图片场景中"
        GenerateImagesRequest generateImagesRequest = GenerateImagesRequest.builder()
                .model("doubao-seedream-4-0-250828")
                .prompt(seedObject.getPrompt())
                .image(seedObject.getImages())
                .responseFormat(ResponseFormat.Url)
                .size("2K")
                .sequentialImageGeneration("disabled")
                .stream(false)
                .watermark(false)
                .build();
        ImagesResponse imagesResponse = service.generateImages(generateImagesRequest);
        service.shutdownExecutor();

        // 遍历所有图片数据
        if (imagesResponse != null && imagesResponse.getData() != null) {
            return imagesResponse.getData().get(0).getUrl();
        }

        throw new Exception("图片生成失败");
    }
}
